#include <stdio.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <unistd.h>
#include <stdlib.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <string.h>
#include <errno.h>
#include <openssl/sha.h> 
#include <unistd.h>  // 用于 sleep 模拟升级

#define BUFFER_SIZE 1024
#define MAX_FILENAME_SIZE 256

// 计算文件的 SHA256 哈希值
int calculate_sha256(const char *filename, unsigned char *hash) {
    FILE *file = fopen(filename, "rb");
    if (!file) {
        perror("打开文件失败");
        return -1;
    }
    
    SHA256_CTX sha256;
    SHA256_Init(&sha256);
    
    unsigned char buffer[BUFFER_SIZE];
    size_t bytes_read;
    
    while ((bytes_read = fread(buffer, 1, BUFFER_SIZE, file))) {
        SHA256_Update(&sha256, buffer, bytes_read);
    }
    
    SHA256_Final(hash, &sha256);
    fclose(file);
    return 0;
}

// 接收文件并验证完整性
int receive_file(int sockfd) {
    char file_info[MAX_FILENAME_SIZE + 20] = {0};
    
    // 1. 接收文件信息（文件名|文件大小）
    ssize_t info_bytes = recv(sockfd, file_info, sizeof(file_info), 0);
    if (info_bytes <= 0) {
        perror("接收文件信息失败");
        return -1;
    }
    
    // 解析文件名和大小
    char *token = strtok(file_info, "|");
    if (!token) {
        fprintf(stderr, "文件信息格式错误\n");
        return -1;
    }
    char filename[MAX_FILENAME_SIZE];
    strncpy(filename, token, sizeof(filename) - 1);
    filename[sizeof(filename) - 1] = '\0';

    token = strtok(NULL, "|");
    if (!token) {
        fprintf(stderr, "文件信息格式错误\n");
        return -1;
    }
    long file_size = atol(token);
    printf("接收升级文件: %s, 大小: %ld 字节\n", filename, file_size);
    
    // 2. 接收文件哈希
    unsigned char expected_hash[SHA256_DIGEST_LENGTH];
    ssize_t hash_bytes = recv(sockfd, expected_hash, sizeof(expected_hash), MSG_WAITALL);
    if (hash_bytes != sizeof(expected_hash)) {
        fprintf(stderr, "接收哈希值失败\n");
        return -1;
    }
    
    // 3. 创建并写入文件内容
    FILE *fp = fopen(filename, "wb");
    if (!fp) {
        perror("创建文件失败");
        return -1;
    }

    char buffer[BUFFER_SIZE];
    long total_received = 0;
    ssize_t bytes_received;
    while (total_received < file_size) {
        bytes_received = recv(sockfd, buffer, BUFFER_SIZE, 0);
        if (bytes_received <= 0) {
            perror(bytes_received ? "接收数据失败" : "连接已关闭");
            fclose(fp);
            return -1;
        }

        size_t bytes_written = fwrite(buffer, 1, bytes_received, fp);
        if (bytes_written != bytes_received) {
            perror("写入文件失败");
            fclose(fp);
            return -1;
        }

        total_received += bytes_received;
        printf("已接收: %.2f%%\r", (float)total_received / file_size * 100);
        fflush(stdout);
    }
    printf("\n文件接收完成，共接收 %ld 字节\n", total_received);
    fclose(fp);
    
    // 4. 验证文件完整性
    if (total_received != file_size) {
        fprintf(stderr, "文件大小不匹配: 期望 %ld, 实际 %ld\n", file_size, total_received);
        return -1;
    }

    unsigned char actual_hash[SHA256_DIGEST_LENGTH];
    if (calculate_sha256(filename, actual_hash) != 0) {
        fprintf(stderr, "计算文件哈希失败\n");
        return -1;
    }
    if (memcmp(expected_hash, actual_hash, SHA256_DIGEST_LENGTH) != 0) {
        fprintf(stderr, "文件完整性验证失败\n");
        return -1;
    }
    printf("文件完整性验证成功\n");
    
    // 5. 发送接收结果
    const char *result = "SUCCESS";
    if (send(sockfd, result, strlen(result), 0) < 0) {
        perror("发送确认失败");
        return -1;
    }
    return 0;
}

// 模拟升级流程
int perform_upgrade(const char *filename) {
    printf("开始执行升级程序: %s\n", filename);
    printf("验证固件签名...\n");
    sleep(1);
    printf("备份当前系统...\n");
    sleep(1);
    printf("写入闪存...\n");
    sleep(2);
    printf("升级完成! 需要重启设备以应用更新\n");
    return 0;
}

int main(int argc, char const *argv[])
{
    // 1. 创建 TCP 套接字
    int tcpfd = socket(AF_INET, SOCK_STREAM, 0);
    if (tcpfd < 0) {
        perror("创建TCP套接字失败");
        return -1;
    }

    // 2. 设置端口复用（避免地址占用）
    int reuse = 1;
    if (setsockopt(tcpfd, SOL_SOCKET, SO_REUSEADDR, &reuse, sizeof(reuse)) < 0) {
        perror("设置SO_REUSEADDR失败");
        close(tcpfd);
        return -1;
    }

    // 3. 绑定套接字到端口 7878
    struct sockaddr_in addr = {0};
    addr.sin_family = AF_INET;
    addr.sin_port = htons(7899);     // 监听端口 
    addr.sin_addr.s_addr = INADDR_ANY; // 允许任意 IP 连接

    if (bind(tcpfd, (struct sockaddr *)&addr, sizeof(addr)) < 0) {
        perror("绑定TCP套接字失败");
        close(tcpfd);
        return -1;
    }

    // 4. 监听连接（最大等待队列 5）
    if (listen(tcpfd, 5) < 0) {
        perror("监听TCP套接字失败");
        close(tcpfd);
        return -1;
    }
    printf("TCP服务器启动，等待客户端连接...\n");

    // 5. 接受客户端连接
    int clifd = accept(tcpfd, NULL, NULL);
    if (clifd < 0) {
        perror("接受TCP连接失败");
        close(tcpfd);
        return -1;
    }
    printf("客户端已连接，准备接收升级文件\n");

    // 6. 接收并升级文件
    if (receive_file(clifd) == 0) {
        printf("升级文件接收成功\n");
        perform_upgrade("*.*"); // 模拟升级（实际可传递动态文件名）
    } else {
        printf("升级文件接收失败\n");
    }

    // 7. 关闭套接字
    close(clifd);
    close(tcpfd);
    return 0;
}